import requests
import gradio as gr
import os
import time

Gitee_AI_API_KEY = os.environ.get(
    "Gitee_AI_API_KEY", "")

API_URL = "https://ai.gitee.com/api/inference/serverless/Q81BJTYV3A14/speech-to-text"
headers = {
    "Authorization": f"Bearer {Gitee_AI_API_KEY}",
    "Content-Type": "audio/flac"
}


def query(filename):
    with open(filename, "rb") as f:
        data = f.read()
    response = requests.post(API_URL, headers=headers, data=data)

    return response


MODEL_NAME = "whisper-large"

FILE_LIMIT_MB = 1000
YT_LENGTH_LIMIT_S = 3600  # limit to 1 hour YouTube files

description = """
使用 Gitee AI Serverless API 实现语音转文本！ [whisper-large](https://ai.gitee.com/serverless-api)
"""


def transcribe(inputs_file_path):
    if inputs_file_path is None:
        raise gr.Error("请录音，并停止录音后提交，或选择音频文件")
    print(inputs_file_path)
    # text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)["text"]
    start_time = time.time()
    res = query(inputs_file_path)
    output = res.json()
    end_time = time.time()
    elapsed_time = end_time - start_time
    print(output)
    if ("text" in output):
        yield [output['text'], elapsed_time]
    else:
        yield [f"API 响应失败! 状态码: {res.status_code}", elapsed_time]


demo = gr.Blocks()

mr_au = gr.Audio(sources="microphone", type="filepath",
                 autoplay=True, label="录音", max_length=60)
mf_transcribe = gr.Interface(
    fn=transcribe,
    inputs=[
        mr_au
    ],
    submit_btn="提交录音",
    clear_btn="清除",
    outputs=[gr.Textbox(label="结果", type="text", show_copy_button=True),
             gr.Textbox(label="耗时", max_lines=1)],
    title="Whisper Large: 语音转文本",
    description=description,
    allow_flagging="never",
    examples=["./ai.wav", "./audio.wav"],
)

au = gr.Audio(sources="upload", type="filepath",
              label="音频文件", autoplay=True, max_length=60)

file_transcribe = gr.Interface(
    fn=transcribe,
    inputs=[au],
    submit_btn="提交文件",
    clear_btn="清除",
    outputs=[gr.Textbox(label="结果", type="text", show_copy_button=True),
             gr.Textbox(label="耗时", max_lines=1)],
    title="Whisper Large: 语音转文本",
    description=description,
    allow_flagging="never",
    examples=["./ai.wav", "./audio.wav"]
)


with demo:
    gr.TabbedInterface([mf_transcribe, file_transcribe], ["录音", "音频文件"])
demo.queue().launch(show_api=False,
                    root_path="https://679s.prod-bd-wx-if.apps.gitee-ai.com")
